# -*- coding: utf-8 -*-
# @Time : 2021/12/29 14:46
# @Author : Huang
# @Email : HuangMJ6016@foxmail.com
# @File : main.py

import os
import networkx as nx

import node2vec

pseudo_count = 0.01


def gen_global_graph(file_path: str) -> nx.DiGraph:
    _graph = nx.DiGraph()
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            parts = line.split("\t\t")
            source = int(parts[0])
            _graph.add_node(source)
            if parts[1] != "null":
                node_freq_strs = parts[1].split("\t")
                # print(len(node_freq_strs))
                for node_freq_str in node_freq_strs:
                    node_freq = node_freq_str.split(":")
                    weight = int(node_freq[1])
                    target = int(node_freq[0])
                    _graph.add_node(target)
                    _graph.add_edge(source, target, global_weight=weight)
    return _graph


def parse_graph(graph_string: str) -> nx.DiGraph:
    graph_string = graph_string.strip()
    parts = graph_string.split('\t')
    edge_strs = parts[4].split()
    res_graph = nx.DiGraph()
    for edge_str in edge_strs:
        source, target, _ = edge_str.split(':')
        source, target = int(source), int(target)
        res_graph.add_node(source)
        res_graph.add_node(target)
        res_graph.add_edge(source, target,
                           edge_weight=global_graph.edges.get((source, target), {'global_weight': 0})[
                                           'global_weight'] + pseudo_count,
                           global_degree=global_graph.out_degree[target] + pseudo_count)
    for edge in res_graph.edges:
        res_graph.add_edge(edge[0], edge[1], local_degree=res_graph.out_degree[edge[1]] + pseudo_count)
    return res_graph


def random_walk(graph: nx.DiGraph, weight_type: str = 'edge_weight'):
    weight_sum, weight_sum_noleaf = 0, 0
    roots, roots_noleaf = list(), list()
    probs, probs_noleaf = list(), list()
    for node, weight in graph.out_degree(weight=weight_type):
        org_weight = weight
        if weight == 0: weight += pseudo_count
        weight_sum += weight
        if org_weight > 0:
            weight_sum_noleaf += weight
        print(node)
    for node, weight in graph.out_degree(weight=weight_type):
        org_weight = weight
        if weight == 0: weight += pseudo_count
        roots.append(node)
        probs.append(weight / weight_sum)
        if org_weight > 0:
            roots_noleaf.append(node)
            probs_noleaf.append(weight / weight_sum_noleaf)
    G = node2vec.Graph(graph, True, 1.0, 1.0, weight_type)
    G.preprocess_transition_probs()


if __name__ == '__main__':
    global_graph_file_path = r'./data/test-net/global_graph.txt'
    global_graph = gen_global_graph(global_graph_file_path)
    read_cas_file = './data/test-net/cascade_train.txt'
    write_cas_file = './data/test-net/random_walks_train.txt'
    with open(read_cas_file, 'r', encoding='utf-8') as rf:
        with open(write_cas_file, 'w', encoding='utf-8') as wf:
            for line in rf:
                cas_graph = parse_graph(line)
                random_walk(cas_graph)
                print(cas_graph)
